import fitz  # PyMuPDF
import requests
import json
import re

# --- 1. 定义大模型服务配置 ---
MODEL_SERVICES = [
    {
        "api_url": "http://18.0.31.1:9981/v1/chat/completions",
        "model": "ds671"
    }
]

# --- 2. 定义脱敏信息生成函数 ---
def mask_name(name):
    """
    根据姓名长度进行脱敏处理。
    """
    if not isinstance(name, str) or not name:
        return "[姓名]"
    name_len = len(name)
    if name_len == 2:
        return f"{name[0]}*"
    elif name_len == 3:
        return f"{name[0]}*{name[2]}"
    elif name_len > 3:
        return f"{name[0]}{'*' * (name_len - 2)}{name[-1]}"
    else:
        return f"{name[0]}*"

def mask_phone_number(phone_number):
    """
    将手机号部分脱敏。
    """
    if not phone_number or len(phone_number) != 11:
        return "[手机号]"
    return phone_number[:3] + "****" + phone_number[7:]

def mask_id_number(id_number):
    """
    将身份证号部分脱敏。
    """
    if not id_number or len(id_number) != 18:
        return "[身份证号]"
    return id_number[:6] + "********" + id_number[14:]

def get_sensitive_info_from_llm(text):
    """
    调用大模型服务，识别并提取敏感信息。
    """
    if not MODEL_SERVICES:
        print("未配置任何大模型服务，无法进行敏感信息识别。")
        return {}

    model_config = MODEL_SERVICES[0]
    api_url = model_config["api_url"]
    model_name = model_config["model"]

    prompt = f"""
    请从以下文本中提取所有姓名、身份证号和手机号。
    如果找到，请以 JSON 格式返回，例如：
    {{
      "names": ["姓名1", "姓名2"],
      "id_numbers": ["身份证号1", "身份证号2"],
      "phone_numbers": ["手机号1", "手机号2"]
    }}
    如果未找到任何信息，则返回空 JSON 对象 {{}}。

    文本内容：
    {text}
    """

    headers = {
        "Content-Type": "application/json"
    }

    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": prompt}],
        "response_format": {"type": "json_object"}
    }

    try:
        response = requests.post(api_url, headers=headers, data=json.dumps(payload), timeout=600)
        response.raise_for_status()

        result_json_str = response.json()["choices"][0]["message"]["content"]
        sensitive_info = json.loads(result_json_str)

        return sensitive_info

    except requests.exceptions.RequestException as e:
        print(f"调用大模型服务时发生网络错误或超时: {e}")
        return {}
    except (json.JSONDecodeError, KeyError) as e:
        print(f"解析大模型返回内容时发生错误: {e}")
        return {}


# --- 3. 整合脱敏逻辑 (动态调整字体大小) ---
def redact_pdf_with_llm_dynamic(input_pdf_path, output_pdf_path):
    """
    结合大模型，对PDF中的姓名、身份证号和手机号进行脱敏处理。
    并根据原始内容动态调整替换文本的字体大小。
    """
    try:
        doc = fitz.open(input_pdf_path)

        full_text_content = ""
        for page_num in range(len(doc)):
            full_text_content += doc.load_page(page_num).get_text() + "\n"

        print("正在调用大模型服务识别敏感信息...")
        all_sensitive_info = get_sensitive_info_from_llm(full_text_content)

        if not any(all_sensitive_info.values()):
            print("未识别到任何敏感信息，无需脱敏。")
            doc.save(output_pdf_path)
            doc.close()
            return

        print(f"大模型识别到以下敏感信息：{all_sensitive_info}")

        redaction_map = {}
        for name in all_sensitive_info.get('names', []):
            redaction_map[name] = mask_name(name)
        for phone in all_sensitive_info.get('phone_numbers', []):
            redaction_map[phone] = mask_phone_number(phone)
        for id_num in all_sensitive_info.get('id_numbers', []):
            redaction_map[id_num] = mask_id_number(id_num)

        for page_num in range(len(doc)):
            page = doc.load_page(page_num)

            for keyword, replacement_text in redaction_map.items():
                # 寻找关键词
                keyword_rects = page.search_for(keyword)

                for rect in keyword_rects:
                    # 获取原始文本块信息以提取字体大小
                    text_blocks = page.get_text("dict")["blocks"]
                    original_fontsize = None
                    for block in text_blocks:
                        if block["type"] == 0:  # 0表示文本块
                            for line in block["lines"]:
                                for span in line["spans"]:
                                    # 检查span的矩形区域是否与关键词矩形重叠
                                    if fitz.Rect(span["bbox"]).intersects(rect):
                                        original_fontsize = span["size"]
                                        break
                                if original_fontsize:
                                    break
                        if original_fontsize:
                            break

                    # 如果没有找到原始字体大小，使用一个默认值
                    if not original_fontsize:
                        original_fontsize = 10

                    # 动态调整字体大小
                    page.add_redact_annot(rect, text=replacement_text, fill=(1, 1, 1), fontname="china-ss", fontsize=original_fontsize)

            page.apply_redactions()

        doc.save(output_pdf_path, garbage=3, deflate=True)
        doc.close()
        print(f"PDF脱敏成功，已保存至：{output_pdf_path}")

    except Exception as e:
        print(f"处理PDF时发生错误: {e}")


# --- 4. 使用示例 ---
input_file = '111.pdf'
output_file = '111_redacted_dynamic_font.pdf'

# 注意：
# 1. 确保你的本地大模型服务已经启动并可访问。
# 2. 确保你的工作目录中存在 '111.pdf' 文件。
redact_pdf_with_llm_dynamic(input_file, output_file)
